{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/llm/yi.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Yi LLMs\n",
    "\n",
    "This notebook shows how to use Yi series LLMs.\n",
    "\n",
    "If you're opening this Notebook on colab, you will need to install LlamaIndex 🦙 and the Yi Python SDK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-yi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fundamental Usage\n",
    "You will need to get an API key from [platform.01.ai](https://platform.01.ai/apikeys). Once you have one, you can either pass it explicity to the model, or use the `YI_API_KEY` environment variable.\n",
    "\n",
    "The details are as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"YI_API_KEY\"] = \"your api key\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call `complete` with a prompt\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The capital of France is Paris.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.llms.yi import Yi\n",
    "\n",
    "llm = Yi(model=\"yi-large\")\n",
    "response = llm.complete(\"What is the capital of France?\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call `chat` with a list of messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: Ahoy, matey! Me name's Captain Blackbeard, but ye can call me Blackbeard for short. Now, what brings ye to me ship? Are ye ready to sail the seven seas in search of treasure and adventure?\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"What is your name\"),\n",
    "]\n",
    "resp = llm.chat(messages)\n",
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming\n",
    "\n",
    "Using `stream_complete` endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.yi import Yi\n",
    "\n",
    "llm = Yi(model=\"yi-large\")\n",
    "response = llm.stream_complete(\"Who is Paul Graham?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham is a British-American computer scientist, entrepreneur, and essayist. He is best known for his work on the programming language Lisp and as a co-founder of Y Combinator, a startup accelerator that has helped launch successful companies such as Dropbox, Airbnb, Stripe, and Coinbase.\n",
      "\n",
      "Graham's career spans several decades and includes founding Viaweb (later sold to Yahoo! and rebranded as Yahoo! Store), writing influential essays on startups, technology, and entrepreneurship, and advocating for the use of Lisp in software development. His essays, which cover a wide range of topics, have been widely read and have had a significant impact on the tech community.\n",
      "\n",
      "Through Y Combinator, Graham has played a pivotal role in the startup ecosystem, providing not only financial support but also mentorship and advice to entrepreneurs. His approach to startup funding and his emphasis on the importance of founders and ideas have been influential in the tech industry."
     ]
    }
   ],
   "source": [
    "for r in response:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using `stream_chat` endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"What is your name\"),\n",
    "]\n",
    "resp = llm.stream_chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As an AI, I don't have a personal name, but you can call me whatever you like! How about Captain AIbeard?"
     ]
    }
   ],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
